from rqalpha.apis import *
import time
from rqalpha.environment import Environment
from rqalpha.my_factors.reform_data import split_factors, read_adj_factor_price
from rqalpha import run_func

config = {
    "base": {
        "start_date": "2021-07-31",
        "end_date": "2021-07-31",
        "frequency": "1d",
        "accounts": {
            "stock": 100000000
        }
    },

    "extra": {
        "log_level": "info",
    },

    "mod": {
        "my_backtest": {
            "enabled": True
        },
        "sys_analyser": {
            "enabled": True,
            "benchmark": "000905.XSHG",
            # "plot": True,
            "plot": False,
            "report_save_path": "C:\\Users\\huajia\Desktop\\rqalpha4\\rqalpha\\plot_result"
        },
        "sys_simulation": {
            "enabled": True,
            "matching_type": "current_bar",
            "volume_limit": False,
            "inactive_limit": True,
            "price_limit": True
        },
        "sys_accounts": {
            "enabled": True,
            "validate_stock_position": False,
            "stock_t1": False
        },
        "incremental": {
            "enabled": True,
            "strategy_id": "1",
            # 是否启用 csv 保存 feeds 功能，可以设置为 MongodbRecorder
            "recorder": "CsvRecorder",
            # 持久化数据输出文件夹
            "persist_folder": "C:\\Users\huajia\\Desktop\\rqalpha4\\rqalpha\\incremental_result",
            # "persist_folder": None,
            # mongodb
            "mongo_url": "mongodb://localhost",
            "mongo_dbname": "rqalpha_records",
        }
    }
}


def run_bt(config, my_defines):
    # 在这个方法中编写任何的初始化逻辑。context对象将会在你的算法策略的任何方法之间做传递。
    def init(context):
        logger.info("init------------")
        _env = Environment.get_instance()
        factor = my_defines['factor_df']
        # print(symbol)
        logger.info("get factors---")
        dt = my_defines['config.base']['start_date']
        # print(dt)
        last_li = my_defines['last_li']
        # print('last li:', last_li)
        q1, q2, q3, q4, q5, q_last = split_factors(factor, dt, last_li)
        my_defines['q1'] = q1
        my_defines['q2'] = q2
        my_defines['q3'] = q3
        my_defines['q4'] = q4
        my_defines['q5'] = q5
        my_defines['q6'] = q_last
        # context.q1 = q1
        # context.q2 = q2
        # context.q3 = q3
        # context.q4 = q4
        # context.q5 = q5
        # context.q6 = q_last
        my_defines['all_queue'].append([dt, *q1])
        my_defines['all_queue'].append([dt, *q2])
        my_defines['all_queue'].append([dt, *q3])
        my_defines['all_queue'].append([dt, *q4])
        my_defines['all_queue'].append([dt, *q5])
        my_defines['all_queue'].append([dt, *q_last])
        my_defines['all_queue'].append([dt, *q1[0:10]])
        # update_universe(context.s1)
        # subscribe(context.s1)

    def open_auction(context, bar_dict):
        print('open_auction ++++++++')
        print(my_defines['is_first_trade'])
        print(len(my_defines['q1']))
        print(my_defines['q1'])
        # print(bar_dict)
        print(len(my_defines['queue1']), my_defines['queue1'])
        print(len(my_defines['q6']), my_defines['q6'])
        # print(len(my_defines['last_li']), my_defines['last_li'])
        # print(context.q1)
        # print(my_defines['all_queue'])
        # time.sleep(500)

        dt = my_defines['config.base']['start_date']
        dt = pd.to_datetime(dt, format='%Y-%m-%d')
        for group in [1]:
            if my_defines['is_first_trade']:
                print('first_run -------- ')
                for order_book_id in my_defines['q{}'.format(group)]:
                    if order_book_id not in my_defines['queue{}'.format(group)]:
                        # order_value(order_book_id, 100000)
                        order_target_percent(order_book_id, 0.01)
                        my_defines['queue{}'.format(group)][order_book_id] = dt
                my_defines['is_first_trade'] = False
            else:
                print('second_run -------- ')
                buy_li = my_defines['q{}'.format(group)][0:10]
                # buy_li = my_defines['q{}'.format(group)][0:2]
                # print(buy_li)
                # print(my_defines['q6'])
                # print(my_defines['queue{}'.format(group)])
                # need_sell_li0 = [i for i in my_defines['queue{}'.format(group)] if i in my_defines['q6']]
                # need_sell_li1 = [i for i in context.portfolio.positions.keys() if i not in
                #                  my_defines['queue{}'.format(group)]]

                li0 = [i for i in buy_li if i in context.portfolio.positions.keys()]
                li1 = [i for i in my_defines['q6'] if i not in li0]
                # need_sell_li = my_defines['q6']
                for order_book_id in li1:
                    if order_book_id in context.portfolio.positions.keys() and \
                            context.portfolio.positions[order_book_id].quantity != 0:
                        price = bar_dict[order_book_id].open
                        order_target_percent(order_book_id, 0, price=price)
                        if order_book_id in my_defines['queue{}'.format(group)]:
                            del my_defines['queue{}'.format(group)][order_book_id]
                for order_book_id in buy_li:
                    # print('2222')
                    if order_book_id in my_defines['queue{}'.format(group)] and dt >\
                            my_defines['queue{}'.format(group)][order_book_id].date():
                        pass
                        # order_shares(order_book_id, -200)
                        # del my_defines['queue{}'.format(group)][order_book_id]
                        # order_shares(order_book_id, 200)
                        # my_defines['queue{}'.format(group)][order_book_id] = dt
                    else:
                        # print('---0', order_book_id)
                        order_target_percent(order_book_id, 0.1)
                        my_defines['queue{}'.format(group)][order_book_id] = dt

    def open_auction2(context, bar_dict):
        print('open_auction ++++++++')
        print(my_defines['is_first_trade'])
        print(len(my_defines['q1']))
        print(my_defines['q1'])
        # print(bar_dict)
        print(len(my_defines['queue1']), my_defines['queue1'])
        print(len(my_defines['q6']), my_defines['q6'])
        dt_str = my_defines['config.base']['start_date']
        dt = pd.to_datetime(dt_str, format='%Y-%m-%d')
        for group in [1]:
            # print(group)
            need_sell_li0 = [i for i in my_defines['queue{}'.format(group)] if i not in my_defines['q{}'.format(group)]]
            need_sell_li1 = [i for i in context.portfolio.positions.keys() if i not in my_defines['queue{}'.format(group)]]
            need_sell_li = list(set(need_sell_li0 + need_sell_li1))
            for order_book_id in need_sell_li:
                # print(bar_dict[order_book_id])
                if order_book_id in context.portfolio.positions.keys() and \
                        context.portfolio.positions[order_book_id].quantity != 0:
                    price = bar_dict[order_book_id].open
                    order_target_percent(order_book_id, 0, price=price)
                    # print('delete------------')
                    if order_book_id in my_defines['queue{}'.format(group)]:
                        del my_defines['queue{}'.format(group)][order_book_id]

            for order_book_id in my_defines['q{}'.format(group)]:
                # print(order_book_id)
                # print(bar_dict[order_book_id].close, bar_dict[order_book_id].open)
                # price = read_adj_factor_price(dt_str, order_book_id)
                # print('get_from adj factor price:', price)
                # time.sleep(500)
                # if price is None:
                #     price = bar_dict[order_book_id].open
                if order_book_id in my_defines['queue{}'.format(group)] and dt > \
                        my_defines['queue{}'.format(group)][order_book_id].date():
                    pass
                    # if context.portfolio.positions[order_book_id].quantity != 0:
                    #     price = bar_dict[order_book_id].open
                    #     order_target_percent(order_book_id, 0, price=price)
                    #     del my_defines['queue{}'.format(group)][order_book_id]
                    # order_target_percent(order_book_id, 0.5)
                    # my_defines['queue{}'.format(group)][order_book_id] = dt
                if order_book_id not in my_defines['queue{}'.format(group)]:
                    order_target_percent(order_book_id, 0.5)
                    my_defines['queue{}'.format(group)][order_book_id] = dt

    def before_trading(context):
        # 在init中定义的context.s1在开启increment开启后会保持一直不变，那么可以在before中定义
        pass

    def handle_bar(context, bar_dict):
        # print(bar_dict['603348.XSHG'])
        # print(bar_dict['601615.XSHG'])
        # print('=================', my_defines['queue1'])
        # time.sleep(500)
        pass

    def after_trading2(context):
        print('after trading----------------')
        # logger.info(data.tail(3))
        context.all_df = pd.DataFrame(data=my_defines['all_queue'])
        context.all_df.to_csv('./context_df.csv')
        # time.sleep(500)

    def after_trading(context):
        print('after trading----------------')
        # logger.info(data.tail(3))
        context.all_df = pd.DataFrame(data=my_defines['all_queue'])
        context.all_df.to_csv('./context_df.csv')
        li = [i for i in my_defines['queue1'] if i not in context.portfolio.positions.keys()]
        for i in li:
            del my_defines['queue1'][i]
        # time.sleep(500)

    run_func(config=config, init=init, open_auction=open_auction, before_trading=before_trading, handle_bar=handle_bar,
             after_trading=after_trading, my_defines=my_defines)


